import numpy as np
import pandas as pd

from scipy.interpolate import make_interp_spline

from paper.trace.tools_coord_trans import latilongalti_to_xyz, xyz_to_rthetaphi, xyz_to_RThetaPhi, xyz_to_latilongalti


######################################## 保存函数 ########################################
def save_2_mat(path_mat, x_list, y_list, R_list, phi_list):
    # 准备保存为指定格式的文件
    n = len(R_list)
    indices = list(range(n))  # 序号: 0,1,2,...
    # 格式化数据（保留12位小数）
    formatted_x = ['{:.12f}'.format(x) for x in x_list]
    formatted_y = ['{:.12f}'.format(y) for y in y_list]
    formatted_R = ['{:.12f}'.format(r) for r in R_list]
    formatted_phi = ['{:.12f}'.format(p) for p in phi_list]
    # 生成三行数据
    line1 = '\t'.join([str(i) for i in indices])  # 序号行
    line2 = '\t'.join(formatted_x)  # x行
    line3 = '\t'.join(formatted_y)  # y行
    line4 = '\t'.join(formatted_R)  # R行
    line5 = '\t'.join(formatted_phi)  # phi行
    # 保存为文本文件（文件名指定为Measure2.mat，但实际是文本格式）
    with open(path_mat, 'w') as f:
        f.write(line1 + '\n')
        f.write(line2 + '\n')
        f.write(line3 + '\n')
        f.write(line4 + '\n')
        f.write(line5 + '\n')


######################################## 插值函数 ########################################
def interpl_xyz_t_1(x_original, y_original, z_original):
    # 原始数据索引作为插值参数
    t = np.linspace(0, 1, len(x_original))  # 归一化参数
    # 创建三次样条插值函数
    spline_x = make_interp_spline(t, x_original, k=3)  # 三次样条（k=3）
    spline_y = make_interp_spline(t, y_original, k=3)
    spline_z = make_interp_spline(t, z_original, k=3)
    # 生成更密集的插值点（增加5倍采样）
    t_new = np.linspace(0, 1, len(x_original) * 5)
    x_smooth = spline_x(t_new)
    y_smooth = spline_y(t_new)
    z_smooth = spline_z(t_new)
    return x_smooth, y_smooth, z_smooth


def interpl_xyz_t_dist(x_original, y_original, z_original):
    """
    改进版：基于累积弦长参数化的三次样条插值，生成更平滑的轨迹
    """
    # 将输入转为 numpy 数组
    x = np.array(x_original)
    y = np.array(y_original)
    z = np.array(z_original)
    n = len(x)

    # === 1. 使用累积弦长作为参数 t ===
    # 计算每一段的距离
    dx = np.diff(x)
    dy = np.diff(y)
    dz = np.diff(z)
    segment_lengths = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)

    # 累积距离作为参数
    t = np.zeros(n)
    t[1:] = np.cumsum(segment_lengths)
    t = t / t[-1]  # 归一化到 [0, 1]

    # === 2. 创建三次样条插值（可选：指定边界条件更平滑）===
    # 可选：使用 'natural' 或 'clamped' 边界条件，这里用默认的 'not-a-knot'
    spline_x = make_interp_spline(t, x, k=3)
    spline_y = make_interp_spline(t, y, k=3)
    spline_z = make_interp_spline(t, z, k=3)

    # === 3. 增加插值密度（例如 10 倍而不是 5 倍）===
    t_new = np.linspace(0, 1, len(x) * 50)  # 提高密度至 10 倍

    x_smooth = spline_x(t_new)
    y_smooth = spline_y(t_new)
    z_smooth = spline_z(t_new)

    return x_smooth, y_smooth, z_smooth


######################################## 主函数 ########################################
def main(path_csv_dataset, path_csv_interp, path_mat):
    # 1.读取csv
    df = pd.read_csv(path_csv_dataset)
    lats, lons, alts = df['latitude'].values, df['longitude'].values, df['altitude'].values
    lat0, lon0, alt0 = lats[0], lons[0], alts[0]
    # 2.(lati,long,alti) --> (x,y,z)
    x, y, z = latilongalti_to_xyz(lats, lons, alts)
    # 3.(x,y,z)插值
    # x_smooth, y_smooth, z_smooth = interpl_xyz_t_1(x, y, z)     # 固定t的三次样条插值
    x_smooth, y_smooth, z_smooth = interpl_xyz_t_dist(x, y, z)     # 固定t的三次样条插值
    # 4.(x,y,z) --> (r,theta,phi), (R,Theta,Phi)
    r_smooth, theta_smooth, phi_smooth = xyz_to_rthetaphi(x_smooth, y_smooth, z_smooth)
    R_smooth, Theta_smooth, Phi_smooth = xyz_to_RThetaPhi(x_smooth, y_smooth, z_smooth)
    lati_smooth, long_smooth, alti_smooth = xyz_to_latilongalti(lat0, lon0, alt0, x_smooth, y_smooth, z_smooth)
    # 5.保存数据
    df_interpl = pd.DataFrame({
        'theta_smooth': theta_smooth,
        'phi_smooth': phi_smooth,
        'r_smooth': r_smooth,
        'lati_smooth': lati_smooth,
        'long_smooth': long_smooth,
        'alti_smooth': alti_smooth,
        'x_smooth': x_smooth,
        'y_smooth': y_smooth,
        'z_smooth': z_smooth,
        'Theta_smooth': Theta_smooth,
        'Phi_smooth': Phi_smooth,
        'R_smooth': R_smooth,
    })
    # 保存全量到.csv
    df_interpl.to_csv(path_csv_interp, index=False)
    print(f"恢复结果已保存到 {path_csv_interp}")
    # 保存到.mat
    save_2_mat(path_mat, x_smooth, y_smooth, R_smooth, Phi_smooth)


def dataset_test_2_mat(path_csv_test, path_mat_test):
    # 1.读取csv
    df = pd.read_csv(path_csv_test)
    x_smooth, y_smooth, R_smooth, Phi_smooth \
        = df['x_smooth'].values, df['y_smooth'].values, df['R_smooth'].values, df['Phi_smooth'].values
    # 2.保存到.mat
    save_2_mat(path_mat_test, x_smooth, y_smooth, R_smooth, Phi_smooth)



if __name__ == "__main__":
    # 读取原始数据集并保存为.csv
    main("../../../../files/archive/uav_navigation_dataset.csv",
         "uav_dataset_interpl_2025-11-11.csv", "UAV_Measure_1_2025-11-11.mat")
    # 将 AE-Conv2DLSTM 测试集转化为 .mat
    # dataset_test_2_mat("../../../../files/archive/uav_dataset_interpl_test_2025-11-11.csv",
    #                    "UAV_Measure_1_test_2025-11-11.mat")
